/*
 * 
 */

package edu.washington.cs.sounddetector.server.sound;

import jAudioFeatureExtractor.ACE.DataTypes.FeatureDefinition;
import jAudioFeatureExtractor.Aggregators.AggregatorContainer;
import jAudioFeatureExtractor.AudioFeatures.FeatureExtractor;
import jAudioFeatureExtractor.jAudioTools.AudioSamples;

import java.io.ByteArrayInputStream;
import java.util.LinkedList;

import javax.sound.sampled.AudioFormat;
import javax.sound.sampled.AudioInputStream;
import javax.sound.sampled.AudioSystem;

import org.apache.commons.lang3.Validate;

/**
 * <p>
 * SoundProcessor is used to extract a set of audio features from a sample of audio. The preferred
 * way to get a SoundProcessor is through the {@link SoundProcessorBuilder}.
 * </p>
 * 
 * <p>
 * This class is a heavily modified version of jAudio 1.0.4's FeatureProcessor.java in
 * src/jAudioFeatureExtractor/jAudioTools/
 * </p>
 * 
 * <p>
 * Definitions:
 * <ul>
 * <li><em>Feature</em> a feature is a number or set of numbers that can be calculated from a set of
 * audio samples, e.g. the <a href="http://en.wikipedia.org/wiki/Spectral_centroid">spectral
 * centroid</a></li>
 * <li><em>Aggregator</em> an aggregator knows how to take a set of features and aggregate them. An
 * example is the mean.</li>
 * </ul>
 * 
 * The SoundProcessor calculates the value(s) of a feature for each window of a sound. An
 * aggregator, such as the mean, takes the value(s) of the feature over each window and combines
 * them into a single value or set of values that represents the whole sound.
 * </p>
 * 
 * <p>
 * Look at jAudioFeatureExtractor.AudioFeatures for a list of features and
 * jAudioFeatureExtractor.Aggregators for a list of aggregators.
 * </p>
 * 
 * <p>
 * SoundProcessor is mutable and not threadsafe.
 * </p>
 * 
 * @author the.dylan.price@gmail.com
 */
public class SoundProcessor
{
    // The window size used for dividing up the recordings to classify.
    private int windowSize;

    // The number of samples that windows are offset by. A value of zero
    // means that there is no window overlap.
    private int windowOverlapOffset;

    // The sampling rate that all recordings are to be converted to before
    // feature extraction.
    private double samplingRate;

    // Whether or not to normalise recordings before feature extraction.
    private boolean normalise;

    // The features that are to be extracted.
    private FeatureExtractor[] featureExtractors;

    // Which features are to be saved after processing. Entries correspond to
    // the featureExtractors field.
    private boolean[] featuresToSave;

    // The dependencies of the features in the feature_extractors field.
    // The first indice corresponds to the feature_extractors indice
    // and the second identifies the number of the dependent feature.
    // The entry identifies the indice of the feature in feature_extractors
    // that corresponds to a dependant feature. The first dimension will be
    // null if there are no dependent features.
    private int[][] featureExtractorDependencies;

    // The longest number of windows of previous features that each feature must
    // have before it can be extracted. The indice corresponds to that of
    // featureExtractors.
    private int[] maxFeatureOffsets;

    private AggregatorContainer container;

    /**
     * Validates and stores the configuration to use for extracting features from audio recordings.
     * 
     * @param windowSize
     *            The size of the windows that the audio recordings are to be broken into, in number
     *            of samples.
     * @param windowOverlap
     *            The fraction of overlap between adjacent windows. Must be between 0.0 and less
     *            than 1.0, with a value of 0.0 meaning no overlap.
     * @param samplingRate
     *            The sampling rate that all recordings are to be converted to before feature
     *            extraction, in Hz.
     * @param normalise
     *            Whether or not to normalise recordings before feature extraction.
     * @param featureExtractors
     *            All features that can be extracted.
     * @param featuresToSave
     *            Which features are to be saved. Entries correspond to the featureExtractors
     *            parameter.
     * @param container
     *            an AggregatorContainer with all the desired aggregators added to it.
     * @throws Exception
     *             Throws an informative exception if the input parameters are invalid.
     */
    public SoundProcessor(int windowSize, double windowOverlap, double samplingRate,
            boolean normalise, FeatureExtractor[] featureExtractors, boolean[] featuresToSave,
            AggregatorContainer container)
    {
        if (container != null)
        {
            if (container.getNumberOfAggregators() == 0)
            {
                throw new IllegalArgumentException(
                        "Saving aggregated values without any aggregators specified");
            }
        }

        if (windowOverlap < 0.0 || windowOverlap >= 1.0)
        {
            throw new IllegalArgumentException("Window overlap fraction is " + windowOverlap
                    + ".\n" + "This value must be 0.0 or above and less than 1.0.");
        }
        if (windowSize < 3)
        {
            throw new IllegalArgumentException("Window size is " + windowSize + ".\n"
                    + "This value must be above 2.");
        }
        boolean oneSelected = false;
        for (int i = 0; i < featuresToSave.length; i++)
        {
            if (featuresToSave[i])
                oneSelected = true;
        }
        if (!oneSelected)
            throw new IllegalArgumentException("No features have been set to be saved.");

        // Save parameters as fields
        this.windowSize = windowSize;
        this.samplingRate = samplingRate;
        this.normalise = normalise;

        // Calculate the window offset
        this.windowOverlapOffset = (int) (windowOverlap * (double) windowSize);

        this.container = container;

        // Find which features need to be extracted and in what order. Also find
        // the indices of dependencies and the maximum offsets for each feature.
        findAndOrderFeaturesToExtract(featureExtractors, featuresToSave);
    }

    /**
     * Fills the feature_extractors, featureExtractorDependencies, maxFeatureOffsets and
     * featuresToSave fields. This involves finding which features need to be extracted and in what
     * order and finding the indices of dependencies and the maximum offsets for each feature.
     * <p>
     * Daniel McEnnis 05-07-05 added feature offset of dependencies to maxOffset
     * 
     * @param all_feature_extractors
     *            All features that can be extracted.
     * @param features_to_save_among_all
     *            Which features are to be saved. Entries correspond to the all_feature_extractors
     *            parameter.
     */
    private void findAndOrderFeaturesToExtract(FeatureExtractor[] all_feature_extractors,
            boolean[] features_to_save_among_all)
    {
        // Find the names of all features
        String[] all_feature_names = new String[all_feature_extractors.length];
        for (int feat = 0; feat < all_feature_extractors.length; feat++)
            all_feature_names[feat] = all_feature_extractors[feat].getFeatureDefinition().name;

        // Find dependencies of all features marked to be extracted.
        // Mark as null if features are not to be extracted. Note that will also
        // be null if there are no dependencies.
        String[][] dependencies = new String[all_feature_extractors.length][];
        for (int feat = 0; feat < all_feature_extractors.length; feat++)
        {
            if (features_to_save_among_all[feat])
                dependencies[feat] = all_feature_extractors[feat].getDepenedencies();
            else
                dependencies[feat] = null;
        }

        // Add dependencies to dependencies and if any features are not marked
        // for
        // saving but are marked as a dependency of a feature that is marked to
        // be
        // saved. Also fill features_to_extract in order to know what features
        // to
        // extract(but not necessarily save).
        boolean done = false;
        boolean[] features_to_extract = new boolean[dependencies.length];
        for (int feat = 0; feat < features_to_extract.length; feat++)
        {
            if (features_to_save_among_all[feat])
                features_to_extract[feat] = true;
            else
                features_to_extract[feat] = false;
        }
        while (!done)
        {
            done = true;
            for (int feat = 0; feat < dependencies.length; feat++)
                if (dependencies[feat] != null)
                    for (int i = 0; i < dependencies[feat].length; i++)
                    {
                        String name = dependencies[feat][i];
                        for (int j = 0; j < all_feature_names.length; j++)
                        {
                            if (name.equals(all_feature_names[j]))
                            {
                                if (!features_to_extract[j])
                                {
                                    features_to_extract[j] = true;
                                    dependencies[j] = all_feature_extractors[j].getDepenedencies();
                                    if (dependencies[j] != null)
                                        done = false;
                                }
                                j = all_feature_names.length;
                            }
                        }
                    }
        }

        // Find the correct order to extract features in by filling the
        // feature_extractors field
        int number_features_to_extract = 0;
        for (int i = 0; i < features_to_extract.length; i++)
            if (features_to_extract[i])
                number_features_to_extract++;
        featureExtractors = new FeatureExtractor[number_features_to_extract];
        featuresToSave = new boolean[number_features_to_extract];
        for (int i = 0; i < featuresToSave.length; i++)
            featuresToSave[i] = false;
        boolean[] feature_added = new boolean[dependencies.length];
        for (int i = 0; i < feature_added.length; i++)
            feature_added[i] = false;
        int current_position = 0;
        done = false;
        while (!done)
        {
            done = true;

            // Add all features that have no remaining dependencies and remove
            // their dependencies from all unadded features
            for (int feat = 0; feat < dependencies.length; feat++)
            {
                if (features_to_extract[feat] && !feature_added[feat])
                    if (dependencies[feat] == null) // add feature if it has no
                    // dependencies
                    {
                        feature_added[feat] = true;
                        featureExtractors[current_position] = all_feature_extractors[feat];
                        featuresToSave[current_position] = features_to_save_among_all[feat];
                        current_position++;
                        done = false;

                        // Remove this dependency from all features that have
                        // it as a dependency and are marked to be extracted
                        for (int i = 0; i < dependencies.length; i++)
                            if (features_to_extract[i] && dependencies[i] != null)
                            {
                                int num_defs = dependencies[i].length;
                                for (int j = 0; j < num_defs; j++)
                                {
                                    if (dependencies[i][j].equals(all_feature_names[feat]))
                                    {
                                        if (dependencies[i].length == 1)
                                        {
                                            dependencies[i] = null;
                                            j = num_defs;
                                        }
                                        else
                                        {
                                            String[] temp = new String[dependencies[i].length - 1];
                                            int m = 0;
                                            for (int k = 0; k < dependencies[i].length; k++)
                                            {
                                                if (k != j)
                                                {
                                                    temp[m] = dependencies[i][k];
                                                    m++;
                                                }
                                            }
                                            dependencies[i] = temp;
                                            j--;
                                            num_defs--;
                                        }
                                    }
                                }
                            }
                    }
            }
        }

        // Find the indices of the feature extractor dependencies for each
        // feature
        // extractor
        featureExtractorDependencies = new int[featureExtractors.length][];
        String[] feature_names = new String[featureExtractors.length];
        for (int feat = 0; feat < feature_names.length; feat++)
        {
            FeatureExtractor featureExtractor = featureExtractors[feat];
            if (featureExtractor != null)
                feature_names[feat] = featureExtractor.getFeatureDefinition().name;
        }
        String[][] feature_dependencies_str = new String[featureExtractors.length][];
        for (int feat = 0; feat < feature_dependencies_str.length; feat++)
            feature_dependencies_str[feat] = featureExtractors[feat].getDepenedencies();
        for (int i = 0; i < feature_dependencies_str.length; i++)
            if (feature_dependencies_str[i] != null)
            {
                featureExtractorDependencies[i] = new int[feature_dependencies_str[i].length];
                for (int j = 0; j < feature_dependencies_str[i].length; j++)
                    for (int k = 0; k < feature_names.length; k++)
                        if (feature_dependencies_str[i][j].equals(feature_names[k]))
                            featureExtractorDependencies[i][j] = k;
            }

        // Find the maximum offset for each feature
        // Daniel McEnnis 5-07-05 added feature offset of dependancies to
        // max_offset
        maxFeatureOffsets = new int[featureExtractors.length];
        for (int i = 0; i < maxFeatureOffsets.length; i++)
        {
            if (featureExtractors[i].getDepenedencyOffsets() == null)
                maxFeatureOffsets[i] = 0;
            else
            {
                int[] these_offsets = featureExtractors[i].getDepenedencyOffsets();
                maxFeatureOffsets[i] = Math.abs(these_offsets[0]
                        + maxFeatureOffsets[featureExtractorDependencies[i][0]]);
                for (int k = 0; k < these_offsets.length; k++)
                {
                    int val = Math.abs(these_offsets[k])
                            + maxFeatureOffsets[featureExtractorDependencies[i][k]];
                    if (val > maxFeatureOffsets[i])
                    {
                        maxFeatureOffsets[i] = val;
                    }
                }
            }
        }
    }

    /**
     * <p>
     * Returns the samples stored in the given byte array.
     * </p>
     * <p>
     * Equivalent to
     * 
     * <pre>
     * sample(new AudioInputStream(new ByteArrayInputStream(sound), format, sound.length), format)
     * </pre>
     * 
     * See {@link #sample(AudioInputStream, AudioFormat)} for more info.
     * </p>
     * 
     * @param sound
     *            the sound to extract samples from
     * @param format
     *            the format of the sound
     * @return The processed audio samples. Values will fall between a minimum of -1 and +1. The
     *         indice identifies the sample number.
     */
    public double[] sample(byte[] sound, AudioFormat format)
    {
        Validate.notNull(sound);
        Validate.notNull(format);

        ByteArrayInputStream bais = new ByteArrayInputStream(sound);
        AudioInputStream stream = new AudioInputStream(bais, format, sound.length);

        return sample(stream, format);
    }

    /**
     * <p>
     * Returns the samples stored in the given AudioInputStream.
     * </p>
     * <p>
     * The samples are re-encoded as PCM signed big endian using the sampling rate in the
     * this.samplingRate field. All channels are projected into one channel. Samples are normalised
     * if this.normalise field is true.
     * </p>
     * 
     * @param stream
     *            The audio input stream to extract samples from.
     * @param format
     *            the format of the audio
     * @return The processed audio samples. Values will fall between a minimum of -1 and +1. The
     *         indice identifies the sample number.
     */
    public double[] sample(AudioInputStream stream, AudioFormat format)
    {
        Validate.notNull(stream);
        Validate.notNull(format);

        // Set the bit depth
        int bitDepth = format.getSampleSizeInBits();
        if (bitDepth != 8 && bitDepth != 16)
            bitDepth = 16;

        // If the audio is not PCM signed big endian, then convert it to PCM
        // signed. This is particularly necessary when dealing with MP3s
        if (format.getEncoding() != AudioFormat.Encoding.PCM_SIGNED
                || format.isBigEndian() == false)
        {
            AudioFormat newFormat = new AudioFormat(AudioFormat.Encoding.PCM_SIGNED,
                    format.getSampleRate(), bitDepth, format.getChannels(), format.getChannels()
                            * (bitDepth / 8), format.getSampleRate(), true);
            stream = AudioSystem.getAudioInputStream(newFormat, stream);
        }

        // Convert to the set sampling rate, if it is not already at this
        // sampling rate. Also, convert to an appropriate bit depth if
        // necessary.
        if (format.getSampleRate() != (float) samplingRate
                || bitDepth != format.getSampleSizeInBits())
        {
            AudioFormat new_format = new AudioFormat(AudioFormat.Encoding.PCM_SIGNED,
                    (float) samplingRate, bitDepth, format.getChannels(), format.getChannels()
                            * (bitDepth / 8), format.getSampleRate(), true);
            stream = AudioSystem.getAudioInputStream(new_format, stream);
        }

        // Extract data from the AudioInputStream
        AudioSamples audioData;
        try
        {
            audioData = new AudioSamples(stream, null, false);
        }
        catch (Exception e)
        {
            // as best as I can tell this is the only reason it will fail -- Dylan
            throw new IllegalArgumentException(e);
        }

        // Normalise samples if this option has been requested
        if (normalise)
            audioData.normalizeMixedDownSamples();

        // Return all channels compressed into one
        return audioData.getSamplesMixedDown();
    }

    /**
     * Extracts the desired features from the given samples.
     * 
     * @param samples
     *            the audio samples that were created from a call to
     *            {@link #sample(AudioInputStream, AudioFormat)} or
     *            {@link #sample(byte[], AudioFormat)}.
     * @return a set of feature values and feature definitions.
     * @throws Exception
     */
    public SoundFeatures extractFeatures(double[] samples) throws Exception
    {
        Validate.notNull(samples);

        // Calculate the window start indices
        LinkedList<Integer> windowStartIndicesList = new LinkedList<Integer>();
        int thisStart = 0;
        while (thisStart < samples.length)
        {
            windowStartIndicesList.add(new Integer(thisStart));
            thisStart += windowSize - windowOverlapOffset;
        }
        Integer[] windowStartIndicesI = windowStartIndicesList.toArray(new Integer[1]);
        int[] windowStartIndices = new int[windowStartIndicesI.length];

        for (int i = 0; i < windowStartIndices.length; i++)
            windowStartIndices[i] = windowStartIndicesI[i].intValue();

        // Extract the feature values from the samples
        double[][][] windowFeatureValues = getFeatures(samples, windowStartIndices);

        container.add(featureExtractors, featuresToSave);
        container.aggregate(windowFeatureValues);

        double[][] results = this.container.getResults();
        FeatureDefinition[] defs = container.getFeatureDefinitions();

        SoundFeatures features = new SoundFeatures(results, defs);

        return features;
    }

    /**
     * Breaks the given samples into the appropriate windows and extracts features from each window.
     * 
     * @param samples
     *            The samples to extract features from. Sample values should generally be between -1
     *            and +1.
     * @param window_start_indices
     *            The indices of samples that correspond to where each window should start.
     * @return The extracted feature values for this recording. The first indice identifies the
     *         window, the second identifies the feature and the third identifies the feature value.
     *         The third dimension will be null if the given feature could not be extracted for the
     *         given window.
     * @throws Exception
     *             Throws an exception if a problem occurs.
     */
    private double[][][] getFeatures(double[] samples, int[] window_start_indices) throws Exception
    {
        // The extracted feature values for this recording. The first indice
        // identifies the window, the second identifies the feature and the
        // third identifies the feature value.
        double[][][] results = new double[window_start_indices.length][featureExtractors.length][];

        // Extract features from each window one by one and add save the
        // results.
        // The last window is zero-padded at the end if it falls off the edge of
        // the
        // provided samples.
        for (int win = 0; win < window_start_indices.length; win++)
        {
            // Find the samples in this window and zero-pad if necessary
            double[] window = new double[windowSize];
            int start_sample = window_start_indices[win];
            int end_sample = start_sample + windowSize - 1;
            if (end_sample < samples.length)
                for (int samp = start_sample; samp <= end_sample; samp++)
                    window[samp - start_sample] = samples[samp];
            else
                for (int samp = start_sample; samp <= end_sample; samp++)
                {
                    if (samp < samples.length)
                        window[samp - start_sample] = samples[samp];
                    else
                        window[samp - start_sample] = 0.0;
                }

            // Extract the features one by one
            for (int feat = 0; feat < featureExtractors.length; feat++)
            {
                // Only extract this feature if enough previous information
                // is available to extract this feature
                if (win >= maxFeatureOffsets[feat])
                {
                    // Find the correct feature
                    FeatureExtractor feature = featureExtractors[feat];

                    // Find previously extracted feature values that this
                    // feature
                    // needs
                    double[][] other_feature_values = null;
                    if (featureExtractorDependencies[feat] != null)
                    {
                        other_feature_values = new double[featureExtractorDependencies[feat].length][];
                        for (int i = 0; i < featureExtractorDependencies[feat].length; i++)
                        {
                            int feature_indice = featureExtractorDependencies[feat][i];
                            int offset = feature.getDepenedencyOffsets()[i];
                            other_feature_values[i] = results[win + offset][feature_indice];
                        }
                    }

                    // Store the extracted feature values
                    results[win][feat] = feature.extractFeature(window, samplingRate,
                            other_feature_values);
                }
                else
                    results[win][feat] = null;
            }
        }

        // Return the results
        return results;
    }
}
